/*
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */
package net.javlov.model;

import java.io.BufferedWriter;
import java.io.FileWriter;
import java.io.IOException;
import java.util.Collection;
import java.util.Comparator;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.TreeSet;

import net.javlov.Option;
import net.javlov.State;
import net.javlov.VectorState;
import net.javlov.model.Model.DiscreteTransitionStats;

@SuppressWarnings("rawtypes")
public class ModelExporter {
	
	static String sep = ",",
			zero = "0",
			end = "\n";
	static VectorState terminalState = new VectorState( new double[]{Double.MAX_VALUE}, true );
	
	public static void export(Model m, List<Option> options, Comparator<? super State> comp) {
		HashMap<State, Integer> stateIds = makeStateIds(m, comp);
		
		int numstates = stateIds.size(),
			numact = options.size();
		//2d matrices consisting of <numact> vertically concatenated
		double[][] transVals = new double[numstates*numact][numstates],
				rewVals = new double[numstates*numact][numstates];
				
		exportToArray(m, options, comp, stateIds, transVals, rewVals);
		writeModels(transVals, "trans.csv", rewVals, "rew.csv");
		writeStates(stateIds, "states.txt");
	}
	
	public static void multiExport(List<Model> lstModel, List<Option> options, Comparator<? super State> comp) {
		HashMap<State, Integer> stateIds = makeStateIds(lstModel, comp);
		
		int numstates = stateIds.size(),
			numact = options.size();
		//2d matrices consisting of <numact> vertically concatenated
		double[][][] allTransVals = new double[lstModel.size()][numstates*numact][numstates],
				allRewVals = new double[lstModel.size()][numstates*numact][numstates];
		double[][] transVals, rewVals;
				
		for ( int i = 0; i < lstModel.size(); i++ ) {
			transVals = allTransVals[i];
			rewVals = allRewVals[i];
			exportToArray(lstModel.get(i), options, comp, stateIds, transVals, rewVals);
			writeModels(transVals, "trans" + i + ".csv", rewVals, "rew" + i + ".csv");
		}
		writeStates(stateIds, "states.txt");
	}
	
	public static void exportToArray(Model m, List<Option> options, Comparator<? super State> comp, Map<State, Integer> stateIds,
			double[][] transVals, double[][] rewVals) {	
		int numstates = stateIds.size();
				
		int currAct, sprimeId, sId;
		State s;
		//for ( State st : stateIds.keySet() )
		//	System.out.println(st + " " + st.hashCode() );
				
		for ( Map.Entry<State, Integer> e : stateIds.entrySet() ) {
			s = e.getKey();
			sId = e.getValue() - 1;
			currAct = 0;
			
			for ( Option o : options ) {
				Collection<DiscreteTransitionStats> transSet = m.getTransitions(s, o).values();
				//check if the state is terminal; if so, create self-transitions with exp reward 0
				if ( s.isTerminal() ) {
					transVals[sId + (currAct*numstates)][sId] = 1;
					rewVals[sId + (currAct*numstates)][sId] = 0;
					currAct++;
					continue; //next option
				}
				
				//no known transition for this (s,o) pair
				if ( transSet == null || transSet.isEmpty() ) {
					currAct++;
					continue; //next option
				}

				for ( DiscreteTransitionStats trans :  transSet ) {
					//if sprime is terminal, create transition to the special terminal state, which should always be the last state
					if ( trans.sprime.isTerminal() )
						sprimeId = numstates - 1;
					else
						sprimeId = stateIds.get(trans.sprime) - 1;
					
					transVals[sId + (currAct*numstates)][sprimeId] = trans.transProb;
					rewVals[sId + (currAct*numstates)][sprimeId] = trans.rewardStats.getMean();
				}
				currAct++;
			}
		}
	}
	
	private static void writeModels(double[][] transVals, String transFile,
						double[][] rewVals, String rewFile) {
		String strTrans = "", strRew = "";
		for ( int i = 0; i < transVals.length; i++ ) {
			for ( int j = 0; j < transVals[i].length-1; j++ ) {
				strTrans += transVals[i][j] + sep;
				strRew += rewVals[i][j] + sep;
			}
			strTrans += transVals[i][transVals[i].length-1] + end;
			strRew += rewVals[i][transVals[i].length-1] + end;
		}
		write(strTrans, transFile);
		write(strRew, rewFile);
	}
	
	private static void writeStates(Map<State, Integer> stateIds, String stateFile) {
		String states = "";
		for ( Map.Entry<State, Integer> e : stateIds.entrySet() ) {
			states += e.getKey() + " " + e.getValue() + "\n";
		}
		write(states, stateFile);
	}
	
	private static void write(String toWrite, String file) {
		try {
			BufferedWriter bw = new BufferedWriter( new FileWriter(file) );
			bw.write(toWrite);
			bw.flush();
			bw.close();
		} catch (IOException ie) {
			System.err.println(ie);
		}
	}
	
	private static HashMap<State, Integer> makeStateIds(Model m, Comparator comp) {
		TreeSet<State> states = new TreeSet<State>(comp);
		
		//sort the states according to comparator
		for ( State s : m.revmodel.keySet() ) { //assume all states are reachable
			if ( s.isTerminal() )
				states.add(terminalState);
			else
				states.add(s);
		}
		
		//assign each state a number
		HashMap<State, Integer> stateIds = new HashMap<State, Integer>(states.size() * 2);
		int i = 1;
		for ( State s : states )
			stateIds.put(s, i++);
		return stateIds;
	}
	
	private static HashMap<State, Integer> makeStateIds(List<Model> lstModel, Comparator comp) {
		TreeSet<State> states = new TreeSet<State>(comp);
		
		for ( Model m : lstModel ) {
			//sort the states according to comparator
			for ( State s : m.revmodel.keySet() ) { //assume all states are reachable
				if ( s.isTerminal() )
					states.add(terminalState);
				else
					states.add(s);
			}
		}
		System.out.println(states);
		//assign each state a number
		HashMap<State, Integer> stateIds = new HashMap<State, Integer>(states.size() * 2);
		int i = 1;
		for ( State s : states )
			stateIds.put(s, i++);
		return stateIds;
	}
}
